Register Tokens for Bounded-State Reasoning in Diffusion Language Models
扩散语言模型有界状态推理的寄存器令牌
🧭 扩散语言模型 · 有界状态推理 ★ 深度解读
University of Wisconsin–Madison · Microsoft Research · UC San Diego
用固定位置的寄存器 token 在分块生成间携带连续状态,优于离散文本携带,代码生成提升显著
📌 工作参考价值(3 点)
- 在长期会话、导购或任务型 Agent 中,若上下文预算受限,可借鉴 register token:在固定位置放少量可学习连续 slot,每轮结束后用一次额外 forward 把状态写入 slot,下一轮复用,避免无限拼接历史;训练时随机 mask prompt 并加入全掩码轮次,防止模型直接看原文绕过状态编码。
- 对分块生成长代码、营销文案或批量商品描述的场景,register carry 可替代 full-context,每 chunk 只多一次 forward,论文实验得到 3.4–5.6× wall-clock 加速,代码 pass@1 高于无 carry 的 full-seq SFT。
- RL 微调时,register 可作为可微状态通道:chunked GRPO 的 advantage 能跨 chunk 回传到 register 写入,适合长 horizon 奖励稀疏任务;同时低维 register 可用线性 probe 做状态审计,判断模型是否记住用户目标、任务进度等。